#include <iostream>

#include <vector>
#include <map>
#include <set>
#include <memory>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <lemon/list_graph.h>
#include <lemon/bin_heap.h>
#include <lemon/unionfind.h>
#include <lemon/dim2.h>
#include <lemon/random.h>

#include <cluster.h>
#include <matrix.h>

using namespace lemon;
using namespace std;
using namespace sds;

typedef ListGraph Graph;
GRAPH_TYPEDEFS(Graph);

typedef Graph::NodeMap<int> IntMap;
typedef Graph::NodeMap<dim2::Point<int> > PosMap;
typedef Graph::EdgeMap<double> WeightMap;
typedef Graph::EdgeMap<int> MultMap;

template <typename _Value>
class PointMap {
public:
  typedef dim2::Point<int> Key;
  typedef _Value Value;

  PointMap(int _width, int _height) 
    : width(_width), height(_height), data(_width * _height) {}

  PointMap(int _width, int _height, int _value) 
    : width(_width), height(_height), data(_width * _height, _value) {}

  Value& operator[](const Key& key) {
    return data[key.x + key.y * width];
  }

  const Value& operator[](const Key& key) const {
    return data[key.x + key.y * width];
  }

  void set(const Key& key, const Value& value) {
    operator[](key) = value;
  }

  int getWidth() const { return width; }
  int getHeight() const { return height; }

private:
  int width, height;
  vector<Value> data;
};

void Read(const string& filename, Graph& graph, PosMap& pos, 
	  WeightMap& weight, MultMap& mult, vector<SInteger>& size) {
  SFile file(filename);

  SDataView<unsigned char> source[CH];
  for (int k = 0; k < CH; ++k) {
    SDataset dataset = file.getDataset(k);
    source[k] = dataset.getDataView<unsigned char>();
  }
  size = source[0].getSize();

  SDataView<Node> nodes(size);

  Cluster dist;
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Node node = graph.addNode();
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      dist += point;
      pos[node].y = i;
      pos[node].x = j;
      nodes[i][j] = node;
    }
  }

  Matrix cm(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      cm(i, j) = dist.covariance(i, j);
    }
  }
  Matrix cmi = cm.inverse();

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j + 1 < size[1]; ++j) {
      Edge e = graph.addEdge(nodes[i][j], nodes[i][j + 1]);
      
      Vector dv(CH);
      for (int k = 0; k < CH; ++k) {
	dv(k) = source[k][i][j] - source[k][i][j + 1];
      }
      weight[e] = exp(-1 * dv * cmi * dv);
      mult[e] = 1;
    }
  }
  for (int i = 0; i + 1 < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Edge e = graph.addEdge(nodes[i][j], nodes[i + 1][j]);

      Vector dv(CH);
      for (int k = 0; k < CH; ++k) {
	dv(k) = source[k][i][j] - source[k][i + 1][j];
      }
      weight[e] = exp(-1 * dv * cmi * dv);
      mult[e] = 1;
    }
  }
}

void Write(const string& filename, Graph& graph, PosMap& pos,
	   UnionFind<PointMap<int> >& uf, vector<SInteger>& size) {
  SFile file(filename, SFile::CREATE);

  IntMap color(graph, -2);
  
  IntMap heapMap(graph, -1);
  BinHeap<int, IntMap> heap(heapMap);
  
  map<int, int> segment;
  
  for (NodeIt it(graph); it != INVALID; ++it) {
    if (uf.size(pos[it]) > 1) {
      heap.push(it, countOutArcs(graph, it));
    } else {
      color[it] = 0;
      segment[uf.find(pos[it])] = 0; 
    }
  }

  vector<Node> order;

  while (!heap.empty()) {
    Node node = heap.top();
    heap.pop();
    color[node] = -1;
    order.push_back(node);
    for (IncEdgeIt it(graph, node); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] == -2) {
	heap.decrease(target, heap[target] - 1);
      }
    }
  }

  for (int i = order.size() - 1; i >= 0; --i) {
    set<int> colors;
    for (IncEdgeIt it(graph, order[i]); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] != -1) {
	colors.insert(color[target]);
      }
    }
    for (int j = 1; j <= 6; ++j) {
      if (colors.find(j) == colors.end()) {
	color[order[i]] = j;
	segment[uf.find(pos[order[i]])] = j; 
	break;
      }
    }
  }

  SDataView<char> dataview(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      dataview[i][j] = segment[uf.find(dim2::Point<int>(j, i))];
    }
  }

  SDataset dataset = file.createDataset("segmented", UINT8, size);
  
  dataset.setDataView<char>(dataview);
}


int main(int argc, const char *argv[]) {
  assert(argc == 3);

  Graph graph;

  PosMap pos(graph);
  WeightMap weight(graph);
  MultMap mult(graph);

  vector<SInteger> size;

  Read(argv[1], graph, pos, weight, mult, size);

  PointMap<int> unionfindMap(size[1], size[0], -1);
  UnionFind<PointMap<int> > unionfind(unionfindMap);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      unionfind.insert(dim2::Point<int>(j, i));
    }
  }

  Graph::EdgeMap<int> heapMap(graph, -1);
  BinHeap<double, Graph::EdgeMap<int> > heap(heapMap);

  std::vector<Edge> edges;
  for (EdgeIt it(graph); it != INVALID; ++it) {
    edges.push_back(it);
  }

  for (int i = 0; i < edges.size(); ++i) {
    swap(edges[i], edges[i + rnd[edges.size() - i]]);
  }

  for (int i = 0; i < edges.size(); ++i) {
    heap.push(edges[i], - weight[edges[i]]);
  }

  //  const int STEP =  size[0] * size[1] - 120000;
  const int STEP =  size[0] * size[1] - 130000;

  for (int k = 0; k < STEP; ++k) {
    if ( (k - 1) * 100 / STEP < k * 100 / STEP) {
      cout << k * 100 / STEP << '%' << endl;
    }

    Edge edge = heap.top();

    Node source = graph.u(edge);
    Node target = graph.v(edge);
    
    Node node = graph.addNode();

    pos[node] = pos[source];

    unionfind.join(pos[source], pos[target]);

    heap.erase(edge);
    graph.erase(edge);

    map<Node, pair<int, double> > targets;
    for (IncEdgeIt it(graph, source); it != INVALID; ++it) {
      pair<int, double> val(mult[it], weight[it]);
      targets.insert(make_pair(graph.runningNode(it), val));
      heap.erase(it);
    }
    graph.erase(source);

    for (IncEdgeIt it(graph, target); it != INVALID; ++it) {
      pair<int, double> val(mult[it], weight[it]);
      map<Node, pair<int, double> >::iterator vt =
	targets.find(graph.runningNode(it));
      if (vt != targets.end()) {
	vt->second.first += val.first;
	vt->second.second += val.second;
      } else {
	targets.insert(make_pair(graph.runningNode(it), val));	
      }
      heap.erase(it);
    }
    graph.erase(target);

    std::vector<Edge> incEdges;
    for (map<Node, pair<int, double> >::iterator 
	   it = targets.begin(); it != targets.end(); ++it) {
      Edge edge = graph.addEdge(node, it->first);
      mult[edge] = it->second.first;
      weight[edge] = it->second.second;
      incEdges.push_back(edge);
    }

    for (int i = 0; i < incEdges.size(); ++i) {
      swap(incEdges[i], incEdges[i + rnd[incEdges.size() - i]]);
    }

    for (int i = 0; i < incEdges.size(); ++i) {
      heap.push(incEdges[i], - weight[incEdges[i]] / mult[incEdges[i]]);
    }    

//     if (k % 500 == 0) {
//       std::ostringstream oss;
//       oss << "meanmergesegmentation_" << k << ".hdf";

//       Write(oss.str(), graph, pos, unionfind, size);
//     }
  }

  Write(argv[2], graph, pos, unionfind, size);

  return 0;
}
